import numpy as np
from onnx import TensorProto

INPUT_TYPE_TO_NP_TYPE = {
    'tensor(float)': np.dtype('float32'),
    'tensor(uint8)': np.dtype('uint8'),
    'tensor(int8)': np.dtype('int8'),
    'tensor(uint16)': np.dtype('uint16'),
    'tensor(int16)': np.dtype('int16'),
    'tensor(int32)': np.dtype('int32'),
    'tensor(int64)': np.dtype('int64'),
    'tensor(bool)': np.dtype('bool'),
    'tensor(float16)': np.dtype('float16'),
    'tensor(float64)': np.dtype('float64'),
    'tensor(complex64)': np.dtype('complex64'),
    'tensor(complex128)': np.dtype('complex128'),
    'tensor(uint32)': np.dtype('uint32'),
    'tensor(uint64)': np.dtype('uint64'),
    'tensor(string)': np.dtype(np.object),

    'seq(tensor(complex64))':np.dtype('complex64'),
    'seq(tensor(complex128))':np.dtype('complex128'),
    'seq(tensor(uint8))':np.dtype('uint8'),
    'seq(tensor(int8))':np.dtype('int8'),
    'seq(tensor(uint16))':np.dtype('uint16'),
    'seq(tensor(int16))':np.dtype('int16'),
    'seq(tensor(uint32))':np.dtype('uint32'),
    'seq(tensor(int32))':np.dtype('int32'),
    'seq(tensor(uint64))':np.dtype('uint64'),
    'seq(tensor(int64))':np.dtype('int64'),
    'seq(tensor(float))':np.dtype('float32'),
    'seq(tensor(float16))':np.dtype('float16'),
    'seq(tensor(double))':np.dtype('double'),
    'seq(tensor(bool))':np.dtype('bool'),
    'seq(tensor(string))':np.dtype(np.object)
}

def get_tensor_type(type):
    types = {
        'tensor(float16)' : TensorProto.FLOAT16,
        'tensor(bfloat16)' : TensorProto.BFLOAT16,
        'tensor(float)' : TensorProto.FLOAT,
        'tensor(int8)' : TensorProto.INT8,
        'tensor(int16)' : TensorProto.INT16,
        'tensor(int32)' : TensorProto.INT32,
        'tensor(int64)' : TensorProto.INT64,
        'tensor(uint8)' : TensorProto.UINT8,
        'tensor(uint16)' : TensorProto.UINT16,
        'tensor(uint32)' : TensorProto.UINT32,
        'tensor(uint64)' : TensorProto.UINT64,
        'tensor(bool)' : TensorProto.BOOL,
        'tensor(double)' : TensorProto.DOUBLE,
        'tensor(string)' :TensorProto.STRING,
        'tensor(complex64)': TensorProto.COMPLEX64,
        'tensor(complex128)': TensorProto.COMPLEX128,

        'seq(tensor(complex64))':TensorProto.COMPLEX64,
        'seq(tensor(complex128))':TensorProto.COMPLEX128,
        'seq(tensor(uint8))':TensorProto.UINT8,
        'seq(tensor(int8))':TensorProto.INT8,
        'seq(tensor(uint16))':TensorProto.UINT16,
        'seq(tensor(int16))':TensorProto.INT16,
        'seq(tensor(uint32))':TensorProto.UINT32,
        'seq(tensor(int32))':TensorProto.INT32,
        'seq(tensor(uint64))':TensorProto.UINT64,
        'seq(tensor(int64))':TensorProto.INT64,
        'seq(tensor(float))':TensorProto.FLOAT,
        'seq(tensor(float16))':TensorProto.FLOAT16,
        'seq(tensor(double))':TensorProto.DOUBLE,
        'seq(tensor(bool))':TensorProto.BOOL,
        'seq(tensor(string))':TensorProto.STRING
    }

    return types.get(type, None)